import os
from PIL import Image, ImageChops  
import pyautogui


def compare_images_with_pillow(x,y,key,image1:Image, image2:Image, tolerance:int=30)->float:  
    """  
    Compare two images using the Pillow (PIL) library.  
    
    :param image_path1: Path to the first image file.  
    :param image_path2: Path to the second image file.  
    :param tolerance: How much the images can differ (in a range of 0-255). Lower values mean stricter comparison.  
    :return: True if the images are the same (or very similar), False otherwise.  
    """  
    # Convert images to the same mode (e.g., RGB)  
    img1 = image1.convert("RGB")  
    img2 = image2.convert("RGB")  

    # Compare the images  
    diff = ImageChops.difference(img1, img2)  

    # Get the differing (non-zero) pixels  
    diff_pixels = diff.getdata()  

    # Count the differing pixels  
    num_diff_pixels = sum(pixel[0] > tolerance for pixel in diff_pixels)
    num = len(img1.getdata())
    conf = float(1) - float(num_diff_pixels)/float(num)
    print(f"比较：{x},{y},{key},{num_diff_pixels},{num},{conf}")
    if x == 9 and y == 5 and key ==8:
        debug = 1

    # Return whether the images are the same or not  
    return conf  

class MGrid:
    def __init__(self,name:str,pos:tuple[int],gridSize:tuple[float],capSize:tuple[float],size:tuple[int],confidence:float):
        self.name = name
        self.pos = pos
        self.gridSize = gridSize
        self.capSize = capSize
        self.size = size
        self.confidence = confidence
        self.custom = {}
        filename = f'prj-{self.name}/config.conf'
        with open(filename,encoding='utf-8') as file:
            for line in file:
                if line.startswith('confidence.'):                 
                    # 去除行尾的换行符，并分割字符串
                    kv = line.strip().split('=')                
                    if len(kv) == 2:  
                        key = kv[0].split('confidence.')[1]  
                        value = kv[1]  
                        self.custom[key] = value 

    def Capture(self):
        dir = f'prj-{self.name}/capture/'
        if not os.path.exists(dir):
            os.makedirs(dir)
        img = pyautogui.screenshot(region=
                                   (self.pos[0],self.pos[1],
                                    int(self.gridSize[0] * float(self.size[0])),
                                    int(self.gridSize[1] * float(self.size[1]))))
        img.save(f"{dir}/{self.name}.png")
        for x in range(self.size[0]):
            for y in range(self.size[1]):
                center = (x*self.gridSize[0]+self.gridSize[0]/2,y*self.gridSize[1]+self.gridSize[1]/2)
                box = (center[0]-self.capSize[0]/2,
                       center[1]-self.capSize[1]/2,
                       center[0]+self.capSize[0]/2,
                       center[1]+self.capSize[1]/2)
                subimg = img.crop(box)
                subimg.save(f'{dir}/{x}-{y:02}.png')
    def __LocateCollect(self,x,y,img:Image,collect:dict,region:tuple,confidence:float=0.9)->int:
        for key, value in collect.items():  
            try:
                conf = compare_images_with_pillow(x,y,key,img,value)
                if conf > confidence:
                    return int(key)
            except Exception as e:
                None                
        return None
    def __LocateCollect1(self,x,y,img:Image,collect:dict,confidence:float=0.9)->int:
        for key, value in collect.items():  
            try:
                poss = pyautogui.locate(needleImage=value,haystackImage=img,confidence=confidence)
                print(f'locate: poss_len={len(poss)},key={key},x={x},y={y}')
                if(len(poss) > 0):
                    return int(key)
            except Exception as e:
                print(f'locate except:  key={key},x={x},y={y},e={e}')
        return None
    def __LocateCollect2(self,x,y,collect:dict,region:tuple,confidence:float=0.9)->int:
        for key, value in collect.items():  
            try:
                poss = None
                if key == 'blank':
                    poss = pyautogui.locateOnScreen(image=value,region=region,confidence=0.99)
                elif key in self.custom:
                    poss = pyautogui.locateOnScreen(image=value,region=region,confidence=float(self.custom[key]))
                else:
                    poss = pyautogui.locateOnScreen(image=value,region=region,confidence=confidence)
                print(f'locate: poss_len={len(poss)},key={key},x={x},y={y}')
                if len(poss) > 0:
                    if key == 'blank':
                        return 0
                    else:
                        return int(key)
                else:
                    rgb_value = pyautogui.pixel(region[0]+region[2]/2, region[1]+region[3]/2)
                    if rgb_value == (150,91,27):
                        return 0
            except Exception as e:
                print(f'locate except:  key={key},x={x},y={y},e={e}')
                rx = region[0]+region[2]/2
                ry = region[1]+region[3]/2
                rgb_value = pyautogui.pixel(int(rx), int(ry))
                if rgb_value == (150,91,27):
                    return 0
        return None
    def __LoadCollect(self,dir:str)->tuple[dict,int]:
        collect = {}
        maxid = 0
        for dirpath, dirnames, filenames in os.walk(dir):  
            for filename in filenames:
                filepath = os.path.join(dirpath, filename)                
                image = Image.open(filepath)
                if image != None:
                    id = filename[0:len(filename)-4]
                    collect[id] = image
                    if id != 'blank' and maxid < int(id):
                        maxid = int(id)
        return (collect,maxid)
  
    def Collect(self):
        dir = f'prj-{self.name}/collect/'
        if not os.path.exists(dir):
            os.makedirs(dir)
        (collect,maxid) = self.__LoadCollect(dir)
        img = pyautogui.screenshot(region=
                                   (self.pos[0],self.pos[1],
                                    int(self.gridSize[0] * float(self.size[0])),
                                    int(self.gridSize[1] * float(self.size[1]))))
        # img.save(f"{dir}/{self.name}.png")
        for x in range(self.size[0]):
            for y in range(self.size[1]):
                center = (x*self.gridSize[0]+self.gridSize[0]/2,y*self.gridSize[1]+self.gridSize[1]/2)
                box = (center[0]-self.capSize[0]/2,
                       center[1]-self.capSize[1]/2,
                       center[0]+self.capSize[0]/2,
                       center[1]+self.capSize[1]/2)
                subimg = img.crop(box)
                # rect = (0,0,self.__gridw()-self.gapSize[0],self.__gridh()-self.gapSize[1])
                rect = (int(self.pos[0]+x*self.gridSize[0]),int(self.pos[1]+y*self.gridSize[1]),
                         int(self.gridSize[0]),int(self.gridSize[1]))
                # b = self.__LocateCollect1(x,y,subimg,collect,confidence=0.7)
                key = self.__LocateCollect2(x,y,collect,rect,self.confidence)
                if not key:
                    maxid += 1
                    subimg.save(f'{dir}/{maxid}.png')
                    collect[maxid] = subimg

    def MakeData(self):
        dir = f'prj-{self.name}/collect/'
        if not os.path.exists(dir):
            os.makedirs(dir)
        dir2 = f'prj-{self.name}/verify/'
        if not os.path.exists(dir2):
            os.makedirs(dir2)
        dir3 = f'prj-{self.name}/data/'
        if not os.path.exists(dir3):
            os.makedirs(dir3)
        (collect,maxid) = self.__LoadCollect(dir)
        img = pyautogui.screenshot(region=
                                   (self.pos[0],self.pos[1],
                                    int(self.gridSize[0] * float(self.size[0])),
                                    int(self.gridSize[1] * float(self.size[1]))))
        # img.save(f"{dir}/{self.name}.png")
        with open(f'{dir3}/{self.name}.txt', 'w',encoding='utf8') as f:
            f.write(f'{self.size[0]}x{self.size[1]}\n')
            for y in range(self.size[1]):
                for x in range(self.size[0]):
                    center = (x*self.gridSize[0]+self.gridSize[0]/2,y*self.gridSize[1]+self.gridSize[1]/2)
                    box = (center[0]-self.capSize[0]/2,
                        center[1]-self.capSize[1]/2,
                        center[0]+self.capSize[0]/2,
                        center[1]+self.capSize[1]/2)
                    subimg = img.crop(box)
                    # rect = (0,0,self.__gridw()-self.gapSize[0],self.__gridh()-self.gapSize[1])
                    rect = (int(self.pos[0]+x*self.gridSize[0]),int(self.pos[1]+y*self.gridSize[1]),
                            int(self.gridSize[0]),int(self.gridSize[1]))
                    # b = self.__LocateCollect1(x,y,subimg,collect,confidence=0.7)
                    if x == 4 and y == 12:
                        debug = 1
                    key = self.__LocateCollect2(x,y,collect,rect,self.confidence)
                    if key == None:
                        maxid += 1
                        subimg.save(f'{dir2}/{maxid}-{x}-{y:02}.png')
                        f.write(f'{maxid:02} ')
                        collect[maxid] = subimg
                    else:
                        f.write(f'{key:02} ')
                f.write('\n')


